{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Cassandra Database Tools\n",
    "\n",
    "Apache Cassandra® is a widely used database for storing transactional application data. The introduction of functions and tooling in Large Language Models has opened up some exciting use cases for existing data in Generative AI applications. The Cassandra Database toolkit enables AI engineers to efficiently integrate Agents with Cassandra data, offering the following features: \n",
    " - Fast data access through optimized queries. Most queries should run in single-digit ms or less. \n",
    " - Schema introspection to enhance LLM reasoning capabilities \n",
    " - Compatibility with various Cassandra deployments, including Apache Cassandra®, DataStax Enterprise™, and DataStax Astra™ \n",
    " - Currently, the toolkit is limited to SELECT queries and schema introspection operations. (Safety first)\n",
    "\n",
    "## Quick Start\n",
    " - Install the cassio library\n",
    " - Set environment variables for the Cassandra database you are connecting to\n",
    " - Initialize CassandraDatabase\n",
    " - Pass the tools to your agent with spec.to_tool_list()\n",
    " - Sit back and watch it do all your work for you\n",
    "\n",
    "## Theory of Operation\n",
    "Cassandra Query Language (CQL) is the primary *human-centric* way of interacting with a Cassandra database. While offering some flexibility when generating queries, it requires knowledge of Cassandra data modeling best practices. LLM function calling gives an agent the ability to reason and then choose a tool to satisfy the request. Agents using LLMs should reason using Cassandra-specific logic when choosing the appropriate tool or chain of tools. This reduces the randomness introduced when LLMs are forced to provide a top-down solution. Do you want an LLM to have complete unfettered access to your database? Yeah. Probably not. To accomplish this, we provide a prompt for use when constructing questions for the agent: \n",
    "\n",
    "```json\n",
    "You are an Apache Cassandra expert query analysis bot with the following features \n",
    "and rules:\n",
    " - You will take a question from the end user about finding specific \n",
    "   data in the database.\n",
    " - You will examine the schema of the database and create a query path. \n",
    " - You will provide the user with the correct query to find the data they are looking \n",
    "   for, showing the steps provided by the query path.\n",
    " - You will use best practices for querying Apache Cassandra using partition keys \n",
    "   and clustering columns.\n",
    " - Avoid using ALLOW FILTERING in the query.\n",
    " - The goal is to find a query path, so it may take querying other tables to get \n",
    "   to the final answer. \n",
    "\n",
    "The following is an example of a query path in JSON format:\n",
    "\n",
    " {\n",
    "  \"query_paths\": [\n",
    "    {\n",
    "      \"description\": \"Direct query to users table using email\",\n",
    "      \"steps\": [\n",
    "        {\n",
    "          \"table\": \"user_credentials\",\n",
    "          \"query\": \n",
    "             \"SELECT userid FROM user_credentials WHERE email = 'example@example.com';\"\n",
    "        },\n",
    "        {\n",
    "          \"table\": \"users\",\n",
    "          \"query\": \"SELECT * FROM users WHERE userid = ?;\"\n",
    "        }\n",
    "      ]\n",
    "    }\n",
    "  ]\n",
    "}\n",
    "```\n",
    "\n",
    "## Tools Provided\n",
    "\n",
    "### `cassandra_db_schema`\n",
    "Gathers all schema information for the connected database or a specific schema. Critical for the agent when determining actions. \n",
    "\n",
    "### `cassandra_db_select_table_data`\n",
    "Selects data from a specific keyspace and table. The agent can pass paramaters for a predicate and limits on the number of returned records. \n",
    "\n",
    "### `cassandra_db_query`\n",
    "Experimental alternative to `cassandra_db_select_table_data` which takes a query string completely formed by the agent instead of parameters. *Warning*: This can lead to unusual queries that may not be as performant(or even work). This may be removed in future releases. If it does something cool, we want to know about that too. You never know!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Environment Setup\n",
    "\n",
    "Install the following Python modules:\n",
    "\n",
    "```bash\n",
    "pip install ipykernel python-dotenv cassio llama-index llama-index-agent-openai llama-index-llms-openai llama-index-tools-cassandra\n",
    "```\n",
    "\n",
    "### .env file\n",
    "Connection is via `cassio` using `auto=True` parameter, and the notebook uses OpenAI. You should create a `.env` file accordingly.\n",
    "\n",
    "For Cassandra, set:\n",
    "```bash\n",
    "CASSANDRA_CONTACT_POINTS\n",
    "CASSANDRA_USERNAME\n",
    "CASSANDRA_PASSWORD\n",
    "CASSANDRA_KEYSPACE\n",
    "```\n",
    "\n",
    "For Astra, set:\n",
    "```bash\n",
    "ASTRA_DB_APPLICATION_TOKEN\n",
    "ASTRA_DB_DATABASE_ID\n",
    "ASTRA_DB_KEYSPACE\n",
    "```\n",
    "\n",
    "For example:\n",
    "\n",
    "```bash\n",
    "# Connection to Astra:\n",
    "ASTRA_DB_DATABASE_ID=a1b2c3d4-...\n",
    "ASTRA_DB_APPLICATION_TOKEN=AstraCS:...\n",
    "ASTRA_DB_KEYSPACE=notebooks\n",
    "\n",
    "# Also set \n",
    "OPENAI_API_KEY=sk-....\n",
    "```\n",
    "\n",
    "(You may also modify the below code to directly connect with `cassio`.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from dotenv import load_dotenv\n",
    "\n",
    "load_dotenv(override=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Import necessary libraries\n",
    "import os\n",
    "\n",
    "import cassio\n",
    "\n",
    "from llama_index.tools.cassandra.base import CassandraDatabaseToolSpec\n",
    "from llama_index.tools.cassandra.cassandra_database_wrapper import (\n",
    "    CassandraDatabase,\n",
    ")\n",
    "\n",
    "from llama_index.core.agent.workflow import FunctionAgent\n",
    "from llama_index.llms.openai import OpenAI"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Connect to a Cassandra Database"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cassio.init(auto=True)\n",
    "\n",
    "session = cassio.config.resolve_session()\n",
    "if not session:\n",
    "    raise Exception(\n",
    "        \"Check environment configuration or manually configure cassio connection parameters\"\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Test data prep\n",
    "\n",
    "session = cassio.config.resolve_session()\n",
    "\n",
    "session.execute(\"\"\"DROP KEYSPACE IF EXISTS llamaindex_agent_test; \"\"\")\n",
    "\n",
    "session.execute(\n",
    "    \"\"\"\n",
    "CREATE KEYSPACE if not exists llamaindex_agent_test \n",
    "WITH replication = {'class': 'SimpleStrategy', 'replication_factor': 1};\n",
    "\"\"\"\n",
    ")\n",
    "\n",
    "session.execute(\n",
    "    \"\"\"\n",
    "    CREATE TABLE IF NOT EXISTS llamaindex_agent_test.user_credentials (\n",
    "    user_email text PRIMARY KEY,\n",
    "    user_id UUID,\n",
    "    password TEXT\n",
    ");\n",
    "\"\"\"\n",
    ")\n",
    "\n",
    "session.execute(\n",
    "    \"\"\"\n",
    "    CREATE TABLE IF NOT EXISTS llamaindex_agent_test.users (\n",
    "    id UUID PRIMARY KEY,\n",
    "    name TEXT,\n",
    "    email TEXT\n",
    ");\"\"\"\n",
    ")\n",
    "\n",
    "session.execute(\n",
    "    \"\"\"\n",
    "    CREATE TABLE IF NOT EXISTS llamaindex_agent_test.user_videos ( \n",
    "    user_id UUID,\n",
    "    video_id UUID,\n",
    "    title TEXT,\n",
    "    description TEXT,\n",
    "    PRIMARY KEY (user_id, video_id)\n",
    ");\n",
    "\"\"\"\n",
    ")\n",
    "\n",
    "user_id = \"522b1fe2-2e36-4cef-a667-cd4237d08b89\"\n",
    "video_id = \"27066014-bad7-9f58-5a30-f63fe03718f6\"\n",
    "\n",
    "session.execute(\n",
    "    f\"\"\"\n",
    "    INSERT INTO llamaindex_agent_test.user_credentials (user_id, user_email) \n",
    "    VALUES ({user_id}, 'patrick@datastax.com');\n",
    "\"\"\"\n",
    ")\n",
    "\n",
    "session.execute(\n",
    "    f\"\"\"\n",
    "    INSERT INTO llamaindex_agent_test.users (id, name, email) \n",
    "    VALUES ({user_id}, 'Patrick McFadin', 'patrick@datastax.com');\n",
    "\"\"\"\n",
    ")\n",
    "\n",
    "session.execute(\n",
    "    f\"\"\"\n",
    "    INSERT INTO llamaindex_agent_test.user_videos (user_id, video_id, title)\n",
    "    VALUES ({user_id}, {video_id}, 'Use Langflow to Build an LLM Application in 5 Minutes');\n",
    "\"\"\"\n",
    ")\n",
    "\n",
    "session.set_keyspace(\"llamaindex_agent_test\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "cassandra_db_schema\n",
      "cassandra_db_schema(keyspace: str) -> List[llama_index.core.schema.Document]\n",
      "Input to this tool is a keyspace name, output is a table description\n",
      "            of Apache Cassandra tables.\n",
      "            If the query is not correct, an error message will be returned.\n",
      "            If an error is returned, report back to the user that the keyspace\n",
      "            doesn't exist and stop.\n",
      "\n",
      "        Args:\n",
      "            keyspace (str): The name of the keyspace for which to return the schema.\n",
      "\n",
      "        Returns:\n",
      "            List[Document]: A list of Document objects, each containing a table description.\n",
      "        \n",
      "<class 'pydantic.main.cassandra_db_schema'>\n",
      "cassandra_db_select_table_data\n",
      "cassandra_db_select_table_data(keyspace: str, table: str, predicate: str, limit: int) -> List[llama_index.core.schema.Document]\n",
      "Tool for getting data from a table in an Apache Cassandra database.\n",
      "            Use the WHERE clause to specify the predicate for the query that uses the\n",
      "            primary key. A blank predicate will return all rows. Avoid this if possible.\n",
      "            Use the limit to specify the number of rows to return. A blank limit will\n",
      "            return all rows.\n",
      "\n",
      "        Args:\n",
      "            keyspace (str): The name of the keyspace containing the table.\n",
      "            table (str): The name of the table for which to return data.\n",
      "            predicate (str): The predicate for the query that uses the primary key.\n",
      "            limit (int): The maximum number of rows to return.\n",
      "\n",
      "        Returns:\n",
      "            List[Document]: A list of Document objects, each containing a row of data.\n",
      "        \n",
      "<class 'pydantic.main.cassandra_db_select_table_data'>\n"
     ]
    }
   ],
   "source": [
    "# Create a CassandraDatabaseToolSpec object\n",
    "db = CassandraDatabase()\n",
    "\n",
    "spec = CassandraDatabaseToolSpec(db=db)\n",
    "\n",
    "tools = spec.to_tool_list()\n",
    "for tool in tools:\n",
    "    print(tool.metadata.name)\n",
    "    print(tool.metadata.description)\n",
    "    print(tool.metadata.fn_schema)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Choose the LLM that will drive the agent\n",
    "# Only certain models support this\n",
    "llm = OpenAI(model=\"gpt-4-1106-preview\")\n",
    "\n",
    "# Create the Agent with our tools. Verbose will echo the agent's actions\n",
    "agent = FunctionAgent(tools=tools, llm=llm)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Invoking the agent with tools\n",
    "We've created an agent that uses an LLM for reasoning and communication with a tool list for actions, Now we can simply ask questions of the agent and watch it utilize the tools we've given it. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Ask our new agent a series of questions. What how the agent uses tools to get the answers.\n",
    "await agent.run(\"What tables are in the keyspace llamaindex_agent_test?\")\n",
    "await agent.run(\"What is the userid for patrick@datastax.com ?\")\n",
    "await agent.run(\"What videos did user patrick@datastax.com upload?\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "llamaindex",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
